import os
import select
import pickle
import socket
import threading
import traceback

from pyim import logger

logger = logger.get_logger(__package__)


class StopPolling(BaseException):
    pass


class Poll:
    class Cmd:
        ev_register = 1
        ev_unregister = 2
        ev_stop = 3

        def __init__(self, cmd, fd=None, ev=None, handler=None):
            self.cmd = cmd
            self.fd = fd
            self.ev = ev

    def __init__(self):
        self._poll = select.poll()
        self._ev_handlers = {}
        self._pipes = os.pipe()
        self._poll_started = False

    def poll(self):
        self._poll_started = True
        self._poll.register(self._pipes[0], select.POLLIN)
        while True:
            try:
                events = self._poll.poll()
            except KeyboardInterrupt:
                break
            except Exception as e:
                logger.error(str(e))
                continue

            try:
                for fd, ev in events:
                    if fd == self._pipes[0]:
                        if ev & select.POLLIN:
                            data = os.read(self._pipes[0], 8)
                            len = int.from_bytes(data, 'big')
                            data = os.read(self._pipes[0], len)
                            cmd = pickle.loads(data)
                            assert isinstance(cmd, Poll.Cmd)
                            if cmd.cmd == self.Cmd.ev_register:
                                self._poll.register(cmd.fd, cmd.ev)
                            elif cmd.cmd == self.Cmd.ev_unregister:
                                self._poll.unregister(cmd.fd)
                                if cmd.fd in self._ev_handlers:
                                    del self._ev_handlers[cmd.fd]
                            elif cmd.cmd == self.Cmd.ev_stop:
                                logger.info("Poll will be exiting ...")
                                raise StopPolling
                    else:
                        if fd in self._ev_handlers:
                            self._ev_handlers[fd](fd, ev)

            except StopPolling:
                break
            except Exception as e:
                traceback.print_exc()
                continue

        self._poll_started = False

    def _send_cmd(self, cmd):
        bs = pickle.dumps(cmd)
        bytes_to_write = (len(bs)).to_bytes(8, byteorder='big') + bs
        total_len = len(bytes_to_write)
        return os.write(self._pipes[1], bytes_to_write) == total_len

    def stop(self):
        cmd = Poll.Cmd(Poll.Cmd.ev_stop)
        return self._send_cmd(cmd)

    def register(self, fd, ev, handler):
        self._ev_handlers[fd] = handler
        if self._poll_started:
            cmd = Poll.Cmd(Poll.Cmd.ev_register, fd, ev, handler)
            return self._send_cmd(cmd)
        else:
            self._poll.register(fd, ev)
            self._ev_handlers[fd] = handler
            return True

    def unregister(self, fd):
        if self._poll_started:
            cmd = Poll.Cmd(Poll.Cmd.ev_unregister, fd)
            return self._send_cmd(cmd)
        else:
            self._poll.unregister(fd)
            if fd in self._ev_handlers:
                del self._ev_handlers[fd]
            return True


class AsyncIO(Poll):
    def __init__(self):
        super(AsyncIO, self).__init__()
        self._write_queue = dict()

    def add(self, fd):
        self.register(fd, select.POLLIN, self._io_handler)

    def remove(self, fd):
        self.unregister(fd)

    def _io_handler(self, fd, ev):
        try:
            if ev & select.POLLIN:
                self.read(fd)
            elif ev & select.POLLOUT:
                if fd not in self._write_queue:
                    self.register(fd, select.POLLIN, self._io_handler)
                    return

                datas = self._write_queue[fd]
                if not datas:
                    del self._write_queue[fd]
                    self.register(fd, select.POLLIN, self._io_handler)
                else:
                    assert isinstance(datas, list)
                    data = datas.pop()
                    self.write(fd, data)
        except StopPolling:
            raise StopPolling

    def queue(self, fd, data):
        if fd not in self._write_queue:
            self._write_queue[fd] = list()

        data_queue = self._write_queue[fd]
        data_queue.append(data)
        self.register(fd, select.POLLIN | select.POLLOUT, self._io_handler)

    def read(self, fd):
        raise NotImplementedError

    def write(self, fd, data):
        raise NotImplementedError


class AsyncNetwork(AsyncIO):
    def get_socket_by_fd(self, fd):
        raise NotImplementedError

    def is_bound(self, sock):
        raise NotImplementedError

    def on_connect(self, sock, server):
        logger.info("{} connected".format(sock.getpeername()))
        return sock

    def on_disconnect(self, sock):
        logger.info("{} disconnected".format(sock.getpeername()))

    def on_received(self, sock, data):
        logger.info("Received: {} < {}".format(data, sock.getpeername()))

    def on_sent(self, sock, data, sz_send):
        logger.info("Sent({}): {} > {}".format(sz_send, data, sock.getpeername()))

    def setsockopt(self, sock):
        raise NotImplementedError

    def read(self, fd):
        sock = self.get_socket_by_fd(fd)

        if self.is_bound(sock):
            client, _ = sock.accept()
            sock = self.on_connect(client, sock)
            if sock:
                self.setsockopt(sock)
                self.add(sock.fileno())
            else:
                sock.close()
        else:
            data = self.recv_data(self.get_socket_by_fd(fd))
            logger.info("Read bytes: {}".format(data))
            if data is None or len(data) == 0:
                self.close(sock)
            else:
                self.on_received(sock, data)

    def write(self, fd, data):
        sock = self.get_socket_by_fd(fd)
        if self.is_bound(sock):
            raise socket.error
        else:
            szdata = self.send_data(sock, data)
            logger.info("Write bytes: {}".format(data))
            return self.on_sent(sock, data, szdata)

    def recv_data(self, sock):
        data = bytes()
        assert isinstance(sock, socket.socket)
        sz_rcv_buf = sock.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF)

        while True:
            try:
                bs = sock.recv(sz_rcv_buf)
                if not bs:
                    raise ConnectionResetError
                data += bs
                if len(bs) < sz_rcv_buf:
                    break
            except socket.error:
                break
            except ConnectionResetError:
                return None

        return data

    def send_data(self, sock, data):
        assert isinstance(sock, socket.socket)
        offset = 0
        total_write = 0
        sz_snd_buf = sock.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)
        while True:
            try:
                sz_write = sock.send(data[offset:offset + sz_snd_buf])
                if sz_write <= 0:
                    return -1

                total_write += sz_write
                if sz_write < sz_snd_buf:
                    break
                offset += sz_write
            except socket.error:
                break

        return total_write

    def send_in_queue(self, sock, data):
        return self.queue(sock.fileno(), data)

    def close(self, sock):
        logger.info("Close remote peer: {}".format(sock.getpeername()))
        self.on_disconnect(sock)
        self.remove(sock.fileno())
        sock.close()


class Client(threading.Thread, AsyncNetwork):
    def __init__(self,
                 address,
                 *args, **kwargs):

        threading.Thread.__init__(self, *args, **kwargs)
        AsyncNetwork.__init__(self)

        self.daemon = True
        self.address = address
        self._lock = threading.Lock()
        self._write_lock = threading.Lock()
        self._write_result = None
        self._exit = False
        self._auto_connect = True
        self.start()

    @property
    def is_connected(self):
        with self._lock:
            return hasattr(self, "socket")

    @property
    def sock(self):
        with self._lock:
            return getattr(self, "socket")

    def run(self):
        while self._auto_connect:
            _sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            self.setsockopt(_sock)
            try:
                _sock.connect(self.address)
                with self._lock:
                    setattr(self, "socket", _sock)
                self.on_connect(_sock, None)
            except Exception as e:
                traceback.print_exc()
                import time
                time.sleep(1)
                logger.info("Can not connect to {}, try it in 1 second".format(self.address))
                continue

            self.add(_sock.fileno())
            self.poll()
            with self._lock:
                delattr(self, "socket")

    def setsockopt(self, sock):
        pass

    def get_socket_by_fd(self, fd):
        return self.sock

    def is_bound(self, sock):
        return False

    def reconnect(self):
        raise StopPolling

    def on_disconnect(self, client):
        super(Client, self).on_disconnect(client)
        self.reconnect()

    def stop(self):
        self._auto_connect = False
        AsyncIO.stop(self)
        self.join()


class Server(AsyncNetwork):
    def __init__(self):
        AsyncIO.__init__(self)
        self._servers = list()
        self._clients = dict()

    def get_socket_by_fd(self, fd):
        for s in self._servers:
            if s.fileno() == fd:
                return s
            elif fd in self._clients[s.fileno()]:
                return self._clients[s.fileno()][fd]

        return None

    def is_bound(self, sock):
        return sock in self._servers

    def is_client(self, fd):
        for s in self._clients.keys():
            if fd in self._clients[s]:
                return self._clients[s][fd]

    def on_connect(self, sock, server):
        super(Server, self).on_connect(sock, server)
        self._clients[server.fileno()][sock.fileno()] = sock
        return sock

    def on_disconnect(self, sock):
        super(Server, self).on_disconnect(sock)
        for s in self._servers:
            if sock.fileno() in self._clients[s.fileno()]:
                del self._clients[s.fileno()][sock.fileno()]

    def get_sockets(self):
        raise NotImplementedError

    def on_start(self):
        raise NotImplementedError

    def start(self):
        for sock in self.get_sockets():
            self._servers.append(sock)
            self.add(sock.fileno())
            self._clients[sock.fileno()] = dict()

        self.on_start()
        self.poll()


if __name__ == "__main__":
    import sys

    if sys.argv[1] == "server":
        server = Server()
